/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package pkwserver;

import java.io.*;
import java.nio.channels.*;
import java.net.*;
import java.nio.ByteBuffer;
import java.security.KeyStore;
import java.security.SecureRandom;
import java.util.*;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLEngine;
import javax.net.ssl.SSLEngineResult;
import javax.net.ssl.SSLEngineResult.HandshakeStatus;
import javax.net.ssl.SSLSession;
import javax.xml.bind.JAXBContext;
import javax.xml.bind.JAXBElement;
import javax.xml.bind.JAXBException;
import javax.xml.bind.Unmarshaller;
import javax.xml.stream.XMLInputFactory;
import javax.xml.stream.XMLStreamException;
import javax.xml.stream.XMLStreamReader;
import pl.pkw.common.VotingData;

/**
 *
 * @author Pawel
 */
public class PKWServer {

    HashMap<SocketChannel, List<byte[]>> connectionsMap;
    Selector selector;
    SSLEngineResult serverResult;	// results from server's last operation
    SSLEngine serverEngine, clientEngine;
    final int BUF_SIZE = 17000;

    public void run() throws Exception {
        String host = "localhost";
        InetSocketAddress socketAddres = new InetSocketAddress(host, 9999);
        try {

            // SSL stuff
            File pKeyFile = new File("pkwkeystore.jks");
            String pKeyPassword = "bskbsk";
            KeyManagerFactory keyManagerFactory = KeyManagerFactory.getInstance("SunX509");
            KeyStore keyStore = KeyStore.getInstance("JKS");

            InputStream keyInput = new FileInputStream(pKeyFile);
            keyStore.load(keyInput, pKeyPassword.toCharArray());
            keyInput.close();
            keyManagerFactory.init(keyStore, pKeyPassword.toCharArray());

            SSLContext context = SSLContext.getInstance("TLS");

            context.init(keyManagerFactory.getKeyManagers(), null, new SecureRandom());
            serverEngine = context.createSSLEngine("localhost", 9999);
            serverEngine.setUseClientMode(false);
            serverEngine.setNeedClientAuth(false);
            serverEngine.setWantClientAuth(true);

            ServerSocketChannel channel = ServerSocketChannel.open();
            channel.configureBlocking(false);
            selector = Selector.open();
            channel.bind(socketAddres);
            channel.register(selector, SelectionKey.OP_ACCEPT);
            connectionsMap = new HashMap<SocketChannel, List<byte[]>>();

            System.out.println(context.getProtocol());
            System.out.println(context.getProvider());

            while (true) {
                selector.select();
                Iterator keys = selector.selectedKeys().iterator();
                while (keys.hasNext()) {
                    SelectionKey key = (SelectionKey) keys.next();
                    keys.remove();

                    if (!key.isValid()) {
                        continue;
                    }

                    if (key.isAcceptable()) {
                        acceptKey(key);
                    }
                    if (key.isReadable()) {
                        read(key);
                    }
                    if (key.isWritable()) {
                        write(key);
                    }
                }
            }
        } catch (IOException e) {
            e.printStackTrace();
        }

    }

    /*
     * If the result indicates that we have outstanding tasks to do,
     * go ahead and run them in this thread.
     */
    private void runDelegatedTasks(SSLEngineResult result,
            SSLEngine engine) throws Exception {
        System.out.println("Delegated");

        if (result.getHandshakeStatus() == HandshakeStatus.NEED_TASK) {
            Runnable runnable;
            while ((runnable = engine.getDelegatedTask()) != null) {
                runnable.run();
            }
            HandshakeStatus hsStatus = engine.getHandshakeStatus();
            if (hsStatus == HandshakeStatus.NEED_TASK) {
                throw new Exception(
                        "handshake shouldn't need additional tasks");
            }
        }
    }

    private void acceptKey(SelectionKey key) throws IOException {
        ServerSocketChannel serverChannel = (ServerSocketChannel) key.channel();
        SocketChannel channel = serverChannel.accept();
        channel.configureBlocking(false);
        Socket sock = channel.socket();
        connectionsMap.put(channel, new ArrayList<byte[]>());
        channel.register(selector, SelectionKey.OP_READ);

        //serverEngine.beginHandshake();
        try {
            doHandshake(key, channel, serverEngine, ByteBuffer.allocate(BUF_SIZE), ByteBuffer.allocate(BUF_SIZE));
        } catch (Exception e) {
            e.printStackTrace();
        }
        SSLSession ses = serverEngine.getSession();
    }

    private void read(SelectionKey key) throws Exception {
        SocketChannel channel = (SocketChannel) key.channel();
        ByteBuffer buffer = ByteBuffer.allocateDirect(BUF_SIZE);
        int numRead = -1;
        try {
            numRead = channel.read(buffer);
        } catch (IOException e) {
            e.printStackTrace();
        }

        if (numRead == -1) {
            connectionsMap.remove(channel);
            // TODO: może jakieś logowanie o zamknięciu połączenia tutaj
            channel.close();
            key.cancel();
            return;
        }



        System.out.println("Got " + numRead + " bytes");

        buffer.flip();
        //byte[] data = new byte[numRead];

        //System.arraycopy(buffer.array(), 0, data, 0, numRead);

//        /numRead=numRead-80;

        ByteBuffer rawData = ByteBuffer.allocate(BUF_SIZE);
        SSLEngineResult res = serverEngine.unwrap(buffer, rawData);
        System.out.println("read status: " + res.getStatus());
        runDelegatedTasks(res, serverEngine);

        //System.out.println(rawData.array().length + " " + new String(rawData.array(), "UTF-8"));
        parseMessage(key, new String(rawData.array(), "UTF-8").trim());
        /*
        //        /ByteBuffer data = ByteBuffer.allocate(numRead);
        
        //rawData.get(data, 0, data.length);
        
        //rawData.wrap(data, 0, numRead);
        
        ByteArrayInputStream bais = new ByteArrayInputStream(rawData.array());
        //Reader byteArray = new InputStreamReader(bais);            
        
        //XMLInputFactory xmlif = XMLInputFactory.newInstance();
        //xmlif.createXMLStreamReader(bais);
        
        
        
        
        
        //StreamReaderDelegate  srd = new StreamReaderDelegate(xmlif.get);
        
        XMLInputFactory xif = XMLInputFactory.newInstance();
        //FileInputStream fis = new FileInputStream(new String(rawData.array(), "UTF-8"));
        //StringReader s = new StringReader(new String(rawData.array(), "UTF-8"));
        XMLStreamReader xsr = xif.createXMLStreamReader(bais);
        System.out.println("text: " + new String(rawData.array()));
        
        System.out.println("---");
        while(xsr.hasNext())
        {
        xsr.nextTag();
        System.out.println(xsr.getLocalName());
        }
        
        System.out.println("---");
        
        
        xsr.nextTag();
        System.out.println("TagName:" + xsr.getLocalName());
        xsr.require(XMLStreamReader.START_ELEMENT, null, "voter");
        
        
        
        JAXBContext jc = JAXBContext.newInstance(new Class[]{Voter.class});
        
        
        
        Unmarshaller u = jc.createUnmarshaller();
        
        JAXBElement<Voter> obj = u.unmarshal(xsr, Voter.class);
        Voter voter = obj.getValue();
        
        // Voter voter = (Voter) u.unmarshal(xsr, (Class<T>) Voter.class);
        //Voter voter = (Voter) u.unmarshal(new StreamSource(new String(rawData.array(), "UTF-8")));
        xsr.close();
        
        
        System.out.println("voter: " + voter.getsId());
        
        
        
         */
        //parseMessage(key, rawData.toString());

        //marshal
        readVoter();

    }

    private void readVoter() throws JAXBException {
    }

    private void write(SelectionKey key) throws Exception {
        SocketChannel channel = (SocketChannel) key.channel();
        List<byte[]> pendingData = connectionsMap.get(channel);
        Iterator<byte[]> items = pendingData.iterator();
        while (items.hasNext()) {
            byte[] item = items.next();
            items.remove();
            // SSL wrapping
            ByteBuffer bb = ByteBuffer.wrap(item);
            ByteBuffer sslData = ByteBuffer.allocate(BUF_SIZE);
            SSLEngineResult res = serverEngine.wrap(bb, sslData);
            System.out.println("write status: " + res.getStatus());
            runDelegatedTasks(res, serverEngine);
            // end SSL wrapping
            int num = channel.write(sslData);
            System.out.println(num);
        }
        key.interestOps(SelectionKey.OP_WRITE);
    }

    private void parseMessage(SelectionKey key, String string) throws JAXBException, XMLStreamException {
        SocketChannel channel = (SocketChannel) key.channel();
        List<byte[]> pendingData = connectionsMap.get(channel);
        System.out.println("data: " + string);
        pendingData.add(("data: " + string).getBytes());


        System.out.println(string.length());

        /*
        int j=0, i = string.length()-1;
        byte[] b = string.getBytes();
        while(i>-1)
        {
        System.out.println(b[i]);
        
        if(b[i]!=(byte)0)
        {
        System.out.println("break");
        break;
        }
        i--;
        }
        
        while(j<b.length)        
        {
        System.out.println(b[j]);
        System.out.println("0: " + Integer.toString( ( b[i] & 0xff ) + 0x100, 16));
        
        System.out.println("a " + Integer.toHexString(b[j]));
        
        
        
        if(Integer.toHexString(b[j]).equals(0x18))
        {
        System.out.println("-----------a " + Integer.toHexString(b[j]));
        break;
        }
        
        j++;
        }
         */
        ByteArrayInputStream in = new ByteArrayInputStream(string.getBytes());
        XMLStreamReader xsr = XMLInputFactory.newInstance().createXMLStreamReader(in);
        xsr.nextTag();
        System.out.println("TagName:" + xsr.getLocalName());

        if ("votingData".equalsIgnoreCase(xsr.getLocalName())) {
            JAXBContext jc = JAXBContext.newInstance(VotingData.class);
            Unmarshaller u = jc.createUnmarshaller();
            xsr.require(XMLStreamReader.START_ELEMENT, null, "votingData");
            VotingData voter = (VotingData) u.unmarshal(xsr);
            System.out.println("voter: " + voter.getsId());
        } else {
            System.out.println("Nothing to bind!");
        }


        /*String string = null;
        SAXBuilder builder = new SAXBuilder();
        try {
        Document document = builder.build(new StringReader(message));
        Element root = document.getRootElement();
        if(root.getName().equals("msg")) {
        Element content = root.getChild("content");
        if(content.getValue().equals("HELLO")) {
        Element msg = new Element("msg");
        Document new_doc = new Document(msg);
        Element new_content = new Element("content");
        new_content.addContent("ACK");
        msg.addContent(new_content);
        XMLOutputter out = new XMLOutputter();
        
        String helloAck = new XMLMessage(out.outputString(new_doc)).toString();
        
        bufferedwriter.write(helloAck);
        bufferedwriter.flush();
        }
        }
        } catch(JDOMException e) {
        System.out.println("Malformed XML");
        System.out.println(e.getMessage());
        } catch(IOException e) {}
        
         */
    }

    /*
     * from http://docs.oracle.com/javase/6/docs/technotes/guides/security/jsse/JSSERefGuide.html
     */
    void doHandshake(SelectionKey key, SocketChannel socketChannel, SSLEngine engine,
            ByteBuffer myNetData, ByteBuffer peerNetData) throws Exception {

        // Create byte buffers to use for holding application data
        int appBufferSize = engine.getSession().getApplicationBufferSize();
        System.out.println("Buffer size: " + appBufferSize);
        ByteBuffer myAppData = ByteBuffer.allocate(appBufferSize);
        ByteBuffer peerAppData = ByteBuffer.allocate(appBufferSize);
        ByteBuffer dummyData = ByteBuffer.allocate(0);

        // Begin handshake
        engine.beginHandshake();
        SSLEngineResult.HandshakeStatus hs = engine.getHandshakeStatus();

        // Process handshaking message
        while (hs != SSLEngineResult.HandshakeStatus.FINISHED
                && hs != SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING) {
            hs = engine.getHandshakeStatus();

            switch (hs) {

                case NEED_UNWRAP:
                    //System.out.println("In NEED_UNWRAP");
                    // Receive handshaking data from peer
                    if (socketChannel.read(peerNetData) < 0) {
                        // Handle closed channel
                    }

                    // Process incoming handshaking data
                    peerNetData.flip();
                    SSLEngineResult res = engine.unwrap(peerNetData, peerAppData);
                    peerNetData.compact();
                    hs = res.getHandshakeStatus();
                    //key.interestOps(SelectionKey.OP_READ);

                    // Check status
                    switch (res.getStatus()) {
                        case OK:
                            // Handle OK status
                            break;

                        // Handle other status: BUFFER_UNDERFLOW, BUFFER_OVERFLOW, CLOSED
                        //...
                    }
                    break;

                case NEED_WRAP:
                    System.out.println("In NEED_WRAP");
                    // Empty the local network packet buffer.
                    myNetData.clear();
                    System.out.println("Limit: " + myNetData.limit());

                    // Generate handshaking data
                    res = engine.wrap(dummyData, myNetData);
                    hs = res.getHandshakeStatus();
                    //myNetData.flip();
                    //socketChannel.write(myNetData);
                    //if(myNetData.hasRemaining())
                    //    key.interestOps(SelectionKey.OP_WRITE);

                    // Check status
                    System.out.println(res.getStatus());
                    switch (res.getStatus()) {
                        case OK:
                            myNetData.flip();

                            // Send the handshaking data to peer
                            while (myNetData.hasRemaining()) {
                                if (socketChannel.write(myNetData) < 0) {
                                    // Handle closed channel
                                }
                            }
                            break;

                        // Handle other status:  BUFFER_OVERFLOW, BUFFER_UNDERFLOW, CLOSED
                        //...
                    }
                    break;

                case NEED_TASK:
                    System.out.println("In NEED_TASK");
                    Runnable r;
                    while ((r = engine.getDelegatedTask()) != null) {
                        System.out.println("Running delegated task");
                        r.run();
                    }
                    break;

                // Handle other status:  // FINISHED or NOT_HANDSHAKING
                //...
            }
        }

        // Processes after handshaking
        //...
    }

    public static void main(String[] args) {
        PKWServer pkw = new PKWServer();
        try {
            pkw.run();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
